{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import tensorflow as tf\n",
    "from keras.layers import *\n",
    "from keras.models import Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def read_tiff(fn):\n",
    "    \"\"\"\n",
    "        inputs: tiff filename\n",
    "        pouputs: image data\n",
    "    \"\"\"\n",
    "    dataset = gdal.Open(fn, gdal.GA_ReadOnly)\n",
    "    im_width = dataset.RasterXSize\n",
    "    im_height = dataset.RasterYSize\n",
    "    im_bands = dataset.RasterCount\n",
    "    im_data = dataset.ReadAsArray(0, 0, im_width, im_height)\n",
    "    if len(im_data.shape) == 2:\n",
    "        return im_data\n",
    "    else:\n",
    "        return im_data.transpose([1, 2, 0])\n",
    "\n",
    "\n",
    "def features(dt):\n",
    "    \"\"\"\n",
    "        inputs: images\n",
    "        pouputs: images with feature bands\n",
    "    \"\"\"\n",
    "    Blue = dt[:, :, 1].reshape(dt.shape[:2] + (1,))\n",
    "    Green = dt[:, :, 2].reshape(dt.shape[:2] + (1,))\n",
    "    Red = dt[:, :, 3].reshape(dt.shape[:2] + (1,))\n",
    "    NIR = dt[:, :, 4].reshape(dt.shape[:2] + (1,))\n",
    "    SWIR1 = dt[:, :, 5].reshape(dt.shape[:2] + (1,))\n",
    "    mndwi = dt[:, :, 5].reshape(dt.shape[:2] + (1,))\n",
    "    leb = dt[:, :, 6].reshape(dt.shape[:2] + (1,))\n",
    "    dem = dt[:, :, 7].reshape(dt.shape[:2] + (1,))\n",
    "    return np.concatenate([Blue, Green, Red, NIR, ndvi, mndwi, leb, dem], axis=-1)\n",
    "\n",
    "\n",
    "def get_input(label, im):\n",
    "    \"\"\"\n",
    "        inputs: None\n",
    "        outputs: model\n",
    "    \"\"\"\n",
    "    features_im = features(im)\n",
    "    y = label[3:-3, 3:-3]\n",
    "    water_pixel_num = (y == 1).sum()\n",
    "    other_pixel_num = (y == 0).sum()\n",
    "    if water_pixel_num < other_pixel_num:\n",
    "        loc_x = np.where(y != 1)[0] + 3\n",
    "        loc_y = np.where(y != 1)[1] + 3\n",
    "        loc_length = len(loc_x)\n",
    "        loc_choice = np.random.choice(np.arange(loc_length), min(int(water_pixel_num * 5), loc_length), replace=False)\n",
    "        loc_x = loc_x[loc_choice]  #np.random.choice(np.where(t!=1)[0],water_pixel_num,replace=False)\n",
    "        loc_y = loc_y[loc_choice]  #np.random.choice(np.where(t!=1)[1],water_pixel_num,replace=False)\n",
    "        label[loc_x, loc_y] = 2\n",
    "    elif water_pixel_num > other_pixel_num:\n",
    "        label[np.where(y == 0)[0] + 3, np.where(y == 0)[1] + 3] = 2\n",
    "        loc_x = np.where(y != 0)[0] + 3\n",
    "        loc_y = np.where(y != 0)[1] + 3\n",
    "        label[loc_x, loc_y] = 0\n",
    "        loc_length = len(loc_x)\n",
    "        loc_choice = np.random.choice(np.arange(loc_length), other_pixel_num, replace=False)\n",
    "        loc_x = loc_x[loc_choice]  #np.random.choice(np.where(t!=1)[0],water_pixel_num,replace=False)\n",
    "        loc_y = loc_y[loc_choice]  #np.random.choice(np.where(t!=1)[1],water_pixel_num,replace=False)\n",
    "        label[loc_x, loc_y] = 1\n",
    "    else:\n",
    "        label = 2 - label\n",
    "    loc_x = np.where(label[3:-3, 3:-3] != 0)[0]\n",
    "    loc_y = np.where(label[3:-3, 3:-3] != 0)[1]\n",
    "    pixel_num = len(loc_x)\n",
    "    #print(pixel_num)\n",
    "    im_train = np.zeros((pixel_num, 7, 7, 8))\n",
    "    lb_train = np.zeros((pixel_num, 1, 1, 2))\n",
    "    num = 0\n",
    "    for i, j in zip(loc_x, loc_y):\n",
    "        try:\n",
    "            im_train[num] = features_im[i:i + 7, j:j + 7, :]\n",
    "        except:\n",
    "            print(i, j)\n",
    "        if label[i + 3, j + 3] == 1:\n",
    "            lb_train[num, :, :, 0] = 1\n",
    "        elif label[i + 3, j + 3] == 2:\n",
    "            lb_train[num, :, :, 1] = 1\n",
    "        num += 1\n",
    "    return im_train, lb_train\n",
    "\n",
    "\n",
    "def CNN():\n",
    "    \"\"\"\n",
    "        inputs: None\n",
    "        outputs: model\n",
    "    \"\"\"\n",
    "    inputs = Input((7, 7, 8))\n",
    "    x = Conv2D(16, (3, 3), activation='relu')(inputs)\n",
    "    x = Conv2D(32, (3, 3), activation='relu')(x)\n",
    "    x = concatenate([x, inputs[:, 2:-2, 2:-2, :]])\n",
    "    x = Conv2D(64, (3, 3), activation='relu')(x)\n",
    "    x = concatenate([x, inputs[:, 3:-3, 3:-3, :]])\n",
    "    x = Conv2D(128, (1, 1), activation='relu')(x)\n",
    "    outputs = Conv2D(2, (1, 1), activation='sigmoid')(x)\n",
    "    model = Model(inputs=inputs, outputs=outputs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def generate_inputs(label_fns, im_fns):\n",
    "    \"\"\"\n",
    "        inputs: filename of labels and images\n",
    "        outputs: None\n",
    "    \"\"\"\n",
    "    for i, j in zip(label_fns, im_fns):\n",
    "        im = read_tiff(j)\n",
    "        label = read_tiff(i)\n",
    "        label = label / label.max()\n",
    "        yield get_input(label, im)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "aa22311dfe557d93"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model = CNN()\n",
    "epochs = {}\n",
    "for ep in range(10):\n",
    "    print('epoch:', ep + 1)\n",
    "    epochs[str(ep + 1)] = {'loss': [], 'acc': []}\n",
    "    i = tf.random.normal([4, 7, 7, 8], mean=-1, stddev=4)\n",
    "    # 指定矩阵的形状\n",
    "    matrix_shape = [4, 1]\n",
    "    # 生成随机数矩阵\n",
    "    random_matrix = tf.random.uniform(shape=matrix_shape, minval=0, maxval=1)\n",
    "    # 将矩阵中的值转换为二进制\n",
    "    j = tf.cast(random_matrix > 0.5, dtype=tf.int32)\n",
    "\n",
    "    model.compile(optimizer=tf.optimizers.Adam(0.001),\n",
    "                  loss='sparse_categorical_crossentropy',\n",
    "                  metrics=['accuracy'])\n",
    "    history = model.fit(i, j, batch_size=64, epochs=1, shuffle=True)  # i 和 j是随机构造的数据\n",
    "    print(history.history['loss'][0])\n",
    "\n",
    "    model.save(r'test.h5')"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1b98a748d9c30eb3"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
